Open system dynamics of the XXZ model with a Lindblad decay flipping the spin.

# pylint: disable=invalid-name

import matplotlib.pyplot as plt
import numpy as np

import qtealeaves as qtl
from qtealeaves import modeling
from qtealeaves.models import get_xxz_1d

plt.rcParams.update({"text.usetex": True, "font.size": 12, "font.family": "lmodern"})


# pylint: disable-next=too-many-locals
def main(tn_type=5, input_folder=None, output_folder=None, plot=True):
    """
    Main method for the dynamics of the XXZ model as Lindblad master equation.

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    input_folder : str | None, optional
        Input folder. Default to None.

    output_folder : str | None, optional
        Output folder. Default to None.

    plot : bool, optional
        If True, plot the results at the end of the example. Default tot True
    """

    def maska(params):
        ll = params["L"]
        tmp = np.zeros((ll), dtype=bool)
        tmp[2:3] = True
        return tmp

    def maskb(params, maska=maska):
        return np.logical_not(maska(params))

    if input_folder is None:
        input_folder = lambda params: "XXZ1d/oqs/in"
    if output_folder is None:
        output_folder = lambda params: "XXZ1d/oqs/out"

    model, my_ops = get_xxz_1d()

    model += modeling.LocalTerm("n", strength="prep", prefactor=-1, mask=maska)
    model += modeling.LocalTerm("n", strength="prep", prefactor=1, mask=maskb)

    # We need splus as a Lindblad operator
    model += modeling.LindbladTerm("splus", strength="gamma")

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(
        max_iter=7, max_bond_dimension=20
    )
    my_ops["splus"] = np.array([[0, 1], [0, 0.]])
    my_ops["sminus"] = np.array([[0, 0], [1, 0.]])

    my_obs = qtl.observables.TNObservables()
    my_obs += qtl.observables.TNObsLocal("<n>", "n")

    # the oqs_mode=2 evolves as quantum trajectory without any jumps
    quench = qtl.DynamicsQuench([0.05] * 20, time_evolution_mode=2, oqs_mode=2)
    quench["prep"] = lambda tt, params: 0.0
    quench["g"] = lambda tt, params: params["g_t"]
    quench["Jx"] = lambda tt, params: params["Jx_t"]
    quench["Jz"] = lambda tt, params: params["Jz_t"]

    params = [
        {
            "L": 8,
            "g": 0.0,
            "Jx": 1e-5,
            "Jz": 1e-5,
            "prep": 200,
            "gamma": 0.125,
            "Jx_t": 0.5,
            "Jz_t": 1,
            "g_t": 0.0,
            "Quenches": [quench],
            "has_log_file": True,
            "exclude_from_hash" : ["Quenches"],
        }
    ]

    simulation = qtl.QuantumGreenTeaSimulation(
        model,
        my_ops,
        my_conv,
        my_obs,
        tn_type=tn_type,
        folder_name_input=input_folder,
        folder_name_output=output_folder,
        has_log_file=False,
        store_checkpoints=False,
    )

    simulation.run(params, delete_existing_folder=True)
    results = simulation.get_static_obs(params[0])

    print("<n(t=0)>", np.sum(results["<n>"]))

    results = simulation.get_dynamic_obs(params[0])[0]
    pop_norm = [elem["norm"] for elem in results]

    if plot:
        fig = plt.figure()
        axis = fig.add_subplot(111)

        tgrid = np.cumsum(quench.dt_grid)
        axis.plot(tgrid, pop_norm, "c-", label="TTN")
        axis.plot(tgrid, np.exp(-0.125 * tgrid), "k:", label="Theory")

        axis.set_xlabel("\\textbf{Evolution time} $\\tau$")
        axis.set_ylabel("\\textbf{Norm}")

        axis.set_yscale("log")

        plt.legend(loc="upper right")
        pdf_name = "XXZ1d/oqs/XXZ_Excitation"
        plt.savefig(pdf_name + ".pdf")

    print("Norm over time", pop_norm)

    print(
        f"\nExample `{__file__}` ran successfully; "
        + "pdf-plots are saved to XXZ1d/oqs folder; "
        + "no asserts implemented."
    )

    return


if __name__ == "__main__":
    main()